import os
from tools.transmit_receive.receiver import Receiver
from tools.o_gan_module import TrafficModelUser, TrafficModelTrainner, TrafficModelParam
from tools.transmit_receive.traffic_template.generator import RandomTimeSpanGenerator
from tools.transmit_receive.traffic_template import TimeInterferencer

EPOCHS = 100
BATCH_SIZE = 64
NOISE_DIM = 8
TRAFFIC_DIM = 16
G_PATH = 'model/g_model.h5'
D_PATH = 'model/d_model.h5'


def get_or_create_gan_model(g_path: str, d_path: str) -> TrafficModelUser:
    """
    创建一个模型使用者
    搭配上一个训练器
    :param g_path:
    :param d_path:
    :return:
    """
    param = TrafficModelParam(TRAFFIC_DIM, NOISE_DIM)
    model_trainer = TrafficModelTrainner(EPOCHS, BATCH_SIZE, param)
    return TrafficModelUser(g_path, d_path, param, model_trainer)


if __name__ == '__main__':
    time_interferencer = TimeInterferencer(0.01)
    model_user = get_or_create_gan_model(os.path.join(os.getcwd(), G_PATH),
                                         os.path.join(os.getcwd(), D_PATH))
    z = model_user.create_noise()
    print("z")
    print(z)
    time_sequence = model_user.generate(z)
    # 模拟了信道的影响
    time_interferencer.intervene(time_sequence)
    z_ = model_user.analysis(time_sequence)
    print("z_")
    print(z_)
    pass
